#-------------------------------------------------------------------------------
# Name:        Average energy of IR/UV spectrum of T.Kneiske et al.
# Purpose:
#
# Author:      GAZ
#
# Created:     25.11.2012
# Copyright:   (c) GAZ 2012
# Licence:     <your licence>
#-------------------------------------------------------------------------------

from pylab import *
from scipy import interpolate
from scipy import integrate
import pickle

class AvrIR:
    def __init__(self):
        self.nEIR  = 201
        self.nzIR  = 201
        self.EIR   = zeros(self.nEIR,float)
        self.DensIR= zeros((self.nEIR,self.nzIR),float)
        AAA        = loadtxt('IRKn.dat', unpack='False')
        self.EIR   = AAA[0][:]
        self.zz    = arange(0.,5.025,0.025)
        self.navr  = zeros(self.nzIR,float)
        self.Enavr = zeros(self.nzIR,float)
        self.Eavr  = zeros(self.nzIR,float)
        for i in range(self.nEIR):
            for j in range(self.nzIR):
                self.DensIR[i][j] = AAA[i+1][j]
        print " Initialization is over"


    def Fninteg(self,E):
        return interpolate.splev(E,self.dn)

    def FEinteg(self,E):
        v = interpolate.splev(E,self.dn)
        return E*v

    def CalcAvr(self):
        arrout = zeros((4,self.nzIR),float)
        for j in range(self.nzIR):
            zz3 = (1.0 + self.zz[j])**3
            self.dn = interpolate.splrep(self.EIR, self.DensIR[j])
            self.navr[j]  = zz3*integrate.quad(self.Fninteg,self.EIR[0],self.EIR[200],epsrel=1E-08, limit=200)[0]
            self.Enavr[j] = zz3*integrate.quad(self.FEinteg,self.EIR[0],self.EIR[200],epsrel=1E-08, limit=200)[0]
            self.Eavr[j]  = self.Enavr[j]/self.navr[j]
            arrout[0][j] = self.zz[j]
            arrout[1][j] = self.navr[j]
            arrout[2][j] = self.Enavr[j]
            arrout[3][j] = self.Eavr[j]
            print (' %f %e %e %e') % (self.zz[j], self.navr[j], self.Enavr[j], self.Eavr[j])

        Out = open('AvrIRUV.pcl',"w")
        pickle.dump((self.navr,self.Enavr,self.Eavr),Out)
        Out.close()

        savetxt('Enavr.dat',transpose(arrout),fmt='%e ', delimiter='   ', newline='\n')

    def ReadAvrIR(self):
        In = open('AvrIRUV.pcl',"r")
        (self.navr,self.Enavr,self.Eavr)=pickle.load(In)
        In.close()

    def PlotDensityIRUV(self):
        plot(self.zz,self.navr)
        xlabel('z')
        ylabel(r'n(z), cm$^{-3}$')
        minorticks_on()
        txt='IR/UV spectrum by T. Kneiske et al. (best fit)'
        text(1.1,2.0,txt,color='b',fontsize=14,backgroundcolor="w")
        grid(True)
        show()

    def PlotAvrIR(self):
        plot(self.zz,self.Eavr)
        xlabel('z')
        ylabel('<E>(z), eV')
        minorticks_on()
        txt='IR/UV spectrum by T. Kneiske et al. (best fit)'
        text(0.63,0.0185,txt,color='b',fontsize=14,backgroundcolor="w")
        grid(True)
        show()

    def PlotEnAvrIR(self):
        plot(self.zz,self.Enavr)
        xlabel('z')
        ylabel(r'<E n(z)>, eV cm$^{-3}$')
        minorticks_on()
        txt='IR/UV spectrum by T. Kneiske et al. (best fit)'
        text(1.13,0.0185,txt,color='b',fontsize=14,backgroundcolor="w")
        grid(True)
        show()

# end class AvrIR
###########################################################################################

avrir = AvrIR()
##avrir.CalcAvr()
avrir.ReadAvrIR()
##avrir.PlotAvrIR()
##avrir.PlotEnAvrIR()
avrir.PlotDensityIRUV()
quit()

dn0 = interpolate.splrep(EIR, DensIR[0])
def ftot(E):
    v = interpolate.splev(E, dn0)
    return v




##ntot = integrate.quad(ftot,EIR[0],EIR[200])
##print " ntot = %e " % ntot[0]
##quit()

plt = 2

for i in range(0,210,10):
##for i in range(0,10,10):
    if plt == 0:
        plot(EIR,(1.+zz[i])**3*DensIR[i])
    elif plt == 1:
        plot(EIR,(1.+zz[i])**3*EIR*DensIR[i])
    elif plt == 2:
        plot(EIR,EIR**2*DensIR[i])
##        plot(EIR,(1.+zz[i])**3*EIR**2*DensIR[i])
##        plot(EIR,(1.+zz[i])**3*EIR**2*ftot(EIR)
##        plot(lgEIR,log10((1.+zz[i])**3*EIR**2)+lgDensIR[i])
    else:
        print ' Wrong value of plt'

txtz = r'z=0.25$\times$ i,  i=0,1,..20'

xscale('log')
xlim(1E-3,20.)
yscale('log')
if plt == 0:
    ylim(1E-7,1E4)
    stxty ='dn$_{\mathrm{IR/UV}}$(E)/dE, eV$^{-1}$ cm$^{-3}$'
    text(1.3,1.5E3,'Kneiske et al.',fontsize=16,color='r')
    text(4.5E-3,2.4E-7,txtz,fontsize=14,color='b')
elif plt == 1:
    ylim(1E-6,20.0)
##    ylim(0.,1.)
    stxty ='Edn$_{\mathrm{IR/UV}}$(E)/dE, cm$^{-3}$'
    text(1.3,8.0,'Kneiske et al.',fontsize=16,color='r')
    text(4.5E-3,2.4E-6,txtz,fontsize=14,color='b')
elif plt == 2:
    ylim(1E-5,0.2)
    stxty ='E$^2$dn$_{\mathrm{IR/UV}}$(E)/dE, eV cm$^{-3}$'
    text(1.3,0.1,'Kneiske et al.',fontsize=16,color='r')
    text(4.5E-3,2.4E-5,txtz,fontsize=14,color='b')
else:
    print ' Wrong value of plt'

xlabel('E, eV',fontsize=14)
ylabel(stxty,fontsize=14)

##spl = UnivariateSpline(lgEIR,log10(EIR**2)+lgDensIR[0],k=5,s=1)
##ys = spl(lgEIR)
##
##plot(lgEIR,ys,color='r')

show()

##print EIR
##EIR,DensIR = loadtxt('IRKn.dat',unpack='True')
##for j in range(nzIR):
##    DensIR[j][:] = loadtxt('IRKn.dat',usecols=j+1,unpack='True')

##print AAA[0][0]
##print AAA[1][0]
##print AAA[201][0]

##print DensIR[200][200]
